"""
基于array实现的最大堆（重点在于思路）
"""

from array import array as Array


class MaxHeap:
    def __init__(self, max_size):
        self._data = Array('B', [0 for _ in range(max_size)])
        self._size = 0

    def __len__(self):
        return self._size

    def capacity(self):
        return len(self._data)

    def add(self, item):
        assert self._size < len(self._data), ""
        self._data[self._size] = item
        self._size += 1
        self.__sift_up(self._size - 1)

    def extract(self):
        assert self._size > 0, ""
        item = self._data[0]
        self._data[0] = self._data[self._size - 1]
        self.__sift_down(0)
        return item

    def __sift_up(self, index):
        """
        上浮
        """

        if index > 0:
            parent = index // 2
            if self._data[index] > self._data[parent]:
                parent_item = self._data[parent]
                self._data[parent] = self._data[index]
                self._data[index] = parent_item
                self.__sift_up(parent)

    def __sift_down(self, index):
        """
        下沉
        """

        left = 2 * index + 1
        right = 2 * index + 2
        largest = index

        if left < self._size and self._data[left] >= self._data[largest]:
            largest = left
        elif right < self._size and self._data[right] >= self._data[largest]:
            largest = right

        if largest != index:
            self._data[index], self._data[largest] = self._data[largest], self._data[index]
            self.__sift_down(largest)


def use_max_heap():
    heap = MaxHeap(6)
    heap.add(2)
    heap.add(3)
    heap.add(1)
    print(heap.extract())


if __name__ == '__main__':
    use_max_heap()
